#pragma once

#include <optional>
#include <functional>

#include "helper/iter/Iterator.h"
#include "helper/ControlFlow.h"

namespace helper {
namespace iter {

template<typename T> 
class Take : public Iterator<T>{
public:
  typedef typename Iterator<T>::Item Item;

  Take(Iterator<T> & iter, size_t n) :_iter(iter), _n(n) {}
  ~Take()override{}

  std::optional<Item> next() override {
    if(_n > 0){
      -- _n;
      return _iter.next();
    }
    else{
      return std::nullopt;
    }
  }

  std::optional<Item> nth(size_t n) override{
    //std::cout << "Take::nth" << std::endl;
    if(_n > n){ // take_n > dest:n  => do nth();
      _n -= n + 1;
      return _iter.nth(n); 
    }
    else{  // take_n <= dest:n => None
      if(_n > 0){ //consume the _iter
        _iter.nth(_n -1);
        _n = 0; 
      }
      return std::nullopt; // take_n = dest:n = 0;
    }
  }

  Result<size_t> advanceBy(size_t n) override {
    //std::cout << "Take::advanceBy" << std::endl;
    size_t min = _n < n ? _n : n;
    return _iter.advanceBy(min)
      .template mapOrElse< Result<size_t> >(
        [this](const auto & e) -> Result<size_t>{
          _n -= e;
          return Result<size_t>::Error(e);
        }
        ,[this, min, n](const auto & v) -> Result<size_t> {
          _n -= min;
          return (min < n) ? Result<size_t>::Error(min) : Result<size_t>::OK();
        }
      );
  }
  
  //size_hint
  
private:
  Iterator<T> & _iter;
  size_t _n;
};

}//namespace iter
}//namespace helper

#include "helper/iter/Skip.h"
